import requests
from core.query.prompt.prompt_builder import PromptBuilder
from core.utils.logging_handler import LoggingHandler
from core.model.model_caller import ModelCaller

class UserInputHandler:
    def __init__(self, user_input, api_key=None, base_url=None, local_model_name=None, remote_model_name=None, explanations=None, conversation_history="", use_remote=True):
        self.logger = LoggingHandler().get_logger()
        self.logger.info(f"初始化 UserInputHandler: {user_input}")
        self.user_input = user_input
        self.api_key = api_key
        self.base_url = base_url
        self.local_model_name = local_model_name
        self.remote_model_name = remote_model_name
        self.explanations = explanations
        self.conversation_history = conversation_history
        self.use_remote = use_remote

        # 初始化 ModelCaller
        self.model_caller = ModelCaller(
            api_key=self.api_key,
            base_url=self.base_url,
            local_model_name=self.local_model_name,
            remote_model_name=self.remote_model_name
        )

    def parse_query(self, df=None, custom_prompt=None):
        self.logger.info("开始解析用户查询")
        
        if custom_prompt:
            prompt = custom_prompt
        else:
            # 使用 PromptBuilder 构建提示词
            prompt_builder = PromptBuilder(df, self.explanations, self.conversation_history, self.user_input)
            prompt = prompt_builder.build_prompt()

        # 使用 ModelCaller 调用模型（本地或远程）
        try:
            response = self.model_caller.call_model(prompt, use_remote=self.use_remote)
            self.logger.info("模型响应解析成功")
            
            # 根据响应内容判断是代码还是自然语言总结
            if response.startswith("code"):
                code_block = response.split('```python')[1].split('```')[0].strip()
                return code_block
            else:
                return response  # 返回自然语言总结

        except Exception as e:
            self.logger.error(f"调用模型API时发生异常: {e}", exc_info=True)
            raise ValueError(f"Error calling model API: {e}")





